import json
import base64
import time
import os, functools
from tqdm import tqdm
import os.path as osp
import shutil
import io, math, xmltodict
from requests import auth
from collections import defaultdict
from numpy.lib.npyio import save
from PIL import Image
from multiprocessing import Pool
import numpy as np
import sys
import cv2
import requests
from tool import filesystem, via_tool

def remove_via_file(data_dir, via_name="via_region_data.rec.json"):
    files = filesystem.get_all_filepath(data_dir, [via_name])
    print(files)
    for f in files:
        os.remove(f)

def rec_to_via2(data_dir, local_save=False, save_no_box=False, via_name="via_region_data.rec.json"):
    sub_dirs = filesystem.get_last_dir(data_dir)

    for sub_dir in sub_dirs:
        if not local_save:
            save_dir = sub_dir + "_two"
            os.makedirs(save_dir)
        else:
            save_dir = sub_dir
        via_path = osp.join(save_dir, via_name)
        if osp.exists(via_path): 
            print("exists. ", via_path)
            continue
        print("deal", sub_dir)

        data_dict = {}
        files = filesystem.get_all_filepath(sub_dir, [".png", ".jpg"])
        print("files:", len(files))
        for file_path in tqdm(files):
            if not osp.exists(file_path): continue
            one_data = {}
            file_size = osp.getsize(file_path)
            filename = osp.basename(file_path)
            one_data["filename"] = filename 
            one_data["size"] = file_size 
            one_data["file_attributes"] = {}

            data = {"img_path": file_path} 
            url = "http://127.0.0.1:9666/api/rec"
            r = requests.post(url, json=data)
            if r.status_code != 200: continue
            rec_data = r.json()
            if rec_data["code"] != 0: 
                print(rec_data)
                continue
            
            if rec_data is None or rec_data["code"] != 0: 
                if save_no_box:
                    rec_data = {"items": []}
                else :
                    continue

            filter_flag = False
            regions = []
            exist_dict = {}
            for item in rec_data["items"]:
                x, y, w, h, label, score = item["x"], item["y"], item["w"], item["h"], item["category"], str(item["score"])
                if exist_dict.get("{}x{}".format(x,y), None): continue
                exist_dict["{}x{}".format(x,y)] = 1

                if w < 24 or h < 24: filter_flag = True

                region = {
                    "shape_attributes": {
                        "name": "rect",
                        "x": x,
                        "y": y,
                        "width": w,
                        "height": h
                    },
                    "region_attributes":{
                        "label": label,
                        "score": score
                    }
                }
                regions.append(region)

            # if (not filter_flag) and len(rec_data["items"]) < 2: continue
            if not local_save:
                shutil.copy(file_path, osp.join(save_dir, filename))

            one_data["regions"] = regions
            data_dict[filename+str(file_size)] = one_data
        
        with open(via_path, "w") as rf:
            rf.write(json.dumps(data_dict))


def rename_img_imp(via_file):

    print(via_file)
    with open(via_file,"r") as rf:
        data_via = json.loads(rf.read())
    
    cur_dir = osp.dirname(via_file)
    new_data = {}
    for img_file in filesystem.get_all_filepath(cur_dir, [".png"]):
        old_name = osp.basename(img_file)
        old_size = osp.getsize(img_file)
        img = cv2.imread(img_file)
        save_path = osp.join(cur_dir, ".".join(old_name.split(".")[:-1])+".jpg")
        cv2.imwrite(save_path, img)
        new_name = osp.basename(save_path)
        new_size = osp.getsize(save_path)

        one_data = data_via.get(old_name+str(old_size), None)
        if one_data is None:
            one_data = {}
            one_data["file_attributes"] = {}
            one_data["regions"] = []
    
        one_data["filename"] = new_name 
        one_data["size"] = new_size 
        new_data[new_name+str(new_size)] = one_data
        os.remove(img_file)

    with open(via_file, "w") as rf:
        rf.write(json.dumps(new_data))


def rename_img(data_dir, via_name="via_region_data.rec.json"):
    files = filesystem.get_all_filepath(data_dir, [via_name])


    thread_count = min(12, len(files)) 
    # print('Parent process %s.' % os.getpid())
    p = Pool(thread_count)
    for via_file in files:
        # xml_to_via_imp(xml_path)
        p.apply_async(rename_img_imp, args=(via_file, ))
    # print('Waiting for all subprocesses done...')
    p.close()
    p.join()


                                

def rename_dir(data_dir):
    last_dirs = filesystem.get_last_dir(data_dir)

    for last_dir in last_dirs:
        print(last_dir)
        if last_dir.endswith("images"):
            dst_dir = osp.join(osp.dirname(last_dir), "imgs")
            shutil.move(last_dir, dst_dir)

if __name__ == "__main__":

    # # data_dir = sys.argv[1]
    # data_dir = "/media/dataset/road_vihicle"
    # local_save = True
    # save_no_box = True
    # rec_to_via2(data_dir, local_save, save_no_box)

    # data_dir = "/home/data/road_vihicle"
    # rename_img(data_dir)

    data_dir = "/home/data/road_vihicle"
    rename_dir(data_dir)

    # via_name="via_region_data.rec.json"
    # remove_via_file(data_dir, via_name)

